use actix::prelude::*;
use futures::future::join_all;
use serde::Deserialize;
use std::collections::HashMap;
use std::fs;
use std::sync::Arc;

// Define messages
struct ExecuteTask {
    input: usize,
    env: Arc<EnvVar>,
}

impl Message for ExecuteTask {
    type Result = Result<usize, ()>;
}

// Define the task actor
struct TaskActor {
    name: String,
    script: String,
}

impl TaskActor {
    fn new(name: &str, script: &str) -> Self {
        TaskActor {
            name: name.to_string(),
            script: script.to_string(),
        }
    }
}

impl Actor for TaskActor {
    type Context = Context<Self>;
}

impl Handler<ExecuteTask> for TaskActor {
    type Result = Result<usize, ()>;

    fn handle(&mut self, _msg: ExecuteTask, _: &mut Context<Self>) -> Self::Result {
        println!("Executing task: {} with script: {}", self.name, self.script);
        Ok(0) // Placeholder for actual execution result
    }
}

// Define the environment variable structure
struct EnvVar {
    vars: HashMap<String, usize>,
}

impl EnvVar {
    fn new() -> Self {
        EnvVar {
            vars: HashMap::new(),
        }
    }

    fn set(&mut self, key: &str, value: usize) {
        self.vars.insert(key.to_string(), value);
    }

    fn get<T: std::str::FromStr>(&self, key: &str) -> Option<T> {
        self.vars.get(key).and_then(|v| v.to_string().parse().ok())
    }
}

// Define a structure to hold the DAG configuration
#[derive(Debug, Deserialize)]
struct TaskConfig {
    name: String,
    depend: Option<Vec<String>>,
    run: RunConfig,
}

#[derive(Debug, Deserialize)]
struct RunConfig {
    #[serde(rename = "type")]
    run_type: String,
    script: String,
}

#[derive(Debug, Deserialize)]
struct DagConfig {
    dagx: HashMap<String, TaskConfig>,
}

// Define the scheduler actor
struct SchedulerActor {
    actors: HashMap<String, Addr<TaskActor>>,
    config: DagConfig,
    env: Arc<EnvVar>,
}

impl SchedulerActor {
    fn new(actors: HashMap<String, Addr<TaskActor>>, config: DagConfig, env: Arc<EnvVar>) -> Self {
        SchedulerActor { actors, config, env }
    }
}

impl Actor for SchedulerActor {
    type Context = Context<Self>;
}

impl Handler<ExecuteTask> for SchedulerActor {
    type Result = ResponseActFuture<Self, Result<usize, ()>>;

    fn handle(&mut self, msg: ExecuteTask, _: &mut Context<Self>) -> Self::Result {
        let futures = self.config.dagx.keys().map(|name| {
            let actor = self.actors.get(name).unwrap().clone();
            actor.send(ExecuteTask {
                input: msg.input,
                env: self.env.clone(),
            })
        });

        Box::pin(
            join_all(futures)
                .into_actor(self)
                .map(|results, _actor, _ctx| {
                    let mut total = 0;
                    for res in results {
                        if let Ok(Ok(val)) = res {
                            total += val;
                        }
                    }
                    Ok(total)
                }),
        )
    }
}

#[actix::main]
async fn main() {
    // Load YAML configuration
    let yaml_str = fs::read_to_string("config/actix_dag_config.yaml").expect("Failed to read YAML file");
    let config: DagConfig = serde_yaml::from_str(&yaml_str).expect("Failed to parse YAML");

    // Create environment
    let mut env = EnvVar::new();
    env.set("base", 2);

    // Create a map to hold task actors
    let mut actors = HashMap::new();

    // Initialize actors based on configuration
    for (name, task_config) in &config.dagx {
        let actor = TaskActor::new(&task_config.name, &task_config.run.script).start();
        actors.insert(name.clone(), actor);
    }

    // Create and start the scheduler actor
    let env = Arc::new(env);
    let scheduler = SchedulerActor::new(actors, config, env.clone()).start();

    // Dispatch tasks
    match scheduler.send(ExecuteTask { input: 0, env }).await {
        Ok(result) => println!("The result is {}.", result.unwrap()),
        Err(_) => println!("Failed to execute DAG."),
    }
}
